package edu.washington.ling.syntric.io;

import java.util.ArrayList;
import java.util.regex.Pattern;

import edu.washington.ling.syntric.struct.tree.Node;
import edu.washington.ling.syntric.struct.tree.Token;
import edu.washington.ling.syntric.struct.tree.Tree;
import edu.washington.ling.syntric.util.Utils;


public class TreeWriterPTB extends TreeWriter {
	
	public String getElemForToken(Token t){
		return t.getWord();
	}
	public String getElemForNode(Node n){
		return String.format("(%s", n.getLabel());
	}
	public ArrayList<String> getTreeElems(Node n){
		ArrayList<String> treeElems = new ArrayList<String>();
		if (n instanceof Token) {
			treeElems.add(getElemForToken((Token) n));
		}
		else {
			treeElems.add(getElemForNode(n));
			for (Node child: n.getPSChildren()) {
				treeElems.addAll(getTreeElems(child));
			}
			String lastElem = treeElems.get(treeElems.size() - 1);
			if (Pattern.matches(".*[^\\s^\\)]\\)$", lastElem)) {
				lastElem = lastElem.concat(" )");
			} else {
				lastElem = lastElem.concat(")");
			}
			treeElems.set(treeElems.size() - 1, lastElem);
		}
		return treeElems;
	}
	@Override
	public TreeRepresentation writeTree(Tree tree) {
		ArrayList<String> lines = new ArrayList<String>();
		ArrayList<String> treeElems = getTreeElems(tree.getPSRoot());
		lines.add(Utils.join(treeElems, " "));
		return new TreeRepresentation(lines, tree.getLayer(), "ptb");
	}

}
